from ete2 import Tree
import random
import numpy
import math
import scipy.special


class RANDOM(object):
    TREE_PATH = '/home/konyushk/workspace/MCTS/Tree/'

    def __init__(self, filename, proximity_type, clustering_type, number_of_images, target):
        self.previous_images = []
        self.number_of_images = number_of_images
        self.tree = Tree(UCT.TREE_PATH + filename + str(self.number_of_images) + '_' + proximity_type + '_' + clustering_type + '.nw')
        self.chosen_node = None
        self.target = target
                
    def ModifyTree(self, images, chosen_image):
        node = self.tree.search_nodes(name = str(chosen_image))[0]
        self.chosen_node = node
        self.previous_images = images[:]
        target_node = self.tree.search_nodes(name = str(self.target))[0]
        nodes_to_target_from_chosen = self.tree.get_distance(target_node, self.chosen_node, topology_only=True)
        return nodes_to_target_from_chosen
            
    def ChooseImages(self, setsize):
        nodes_to_chosen_image = 0
        images = []
        while len(images) < setsize:
            current_node = self.tree.get_tree_root()
            while not current_node.is_leaf():
                max_selection_criteria = 0
                for node in current_node.get_children():
                    selection_criteria = random.random()
                    if selection_criteria > max_selection_criteria:
                        max_selection_criteria = selection_criteria
                        max_node = node
                current_node = max_node
            if int(current_node.name) not in images:
                if int(current_node.name) not in self.previous_images:
                    nodes_to_chosen_image += self.tree.get_distance(current_node, self.chosen_node, topology_only=True)
                    images.append(int(current_node.name))
        avnodes_to_chosen_image = nodes_to_chosen_image/setsize
        return images, avnodes_to_chosen_image


class THOMPSON(object):
    TREE_PATH = '/home/konyushk/workspace/MCTS/Tree/'
    ALPHA = 1
    BETA = 1

    def __init__(self, filename, proximity_type, clustering_type, number_of_images, target):
        self.previous_images = []
        self.number_of_images = number_of_images
        self.tree = Tree(THOMPSON.TREE_PATH + filename + str(self.number_of_images) + '_' + proximity_type + '_' + clustering_type + '.nw')
        self.chosen_node = None
        # Add alpha and beta features to the tree
        for node in self.tree.traverse('postorder'):
            prior_success = 10/float(len(node.get_leaves()))
            #node.add_features(alpha = THOMPSON.ALPHA, beta = THOMPSON.BETA, isUpdated = False)
            node.add_features(alpha = prior_success, beta = THOMPSON.BETA, isUpdated = False)
        self.target = target
        
    def ModifyTree(self, images, chosen_image):
        
        for node in self.tree.traverse('postorder'):
            node.isUpdated = False
        self.previous_images = images[:]
        images.remove(chosen_image)
        not_chosen_images = images
        node = self.tree.search_nodes(name = str(chosen_image))[0]
        self.chosen_node = node
        
        node = node.up
#        print 'Chosen image parameters = ', node.alpha, ', ', node.beta
        while node:
            node.alpha += 1
            node.isUpdated = True
            node = node.up
        for image in not_chosen_images:
            node = self.tree.search_nodes(name = str(image))[0]
            node = node.up
            while node:
                if node.isUpdated == False:
                    node.beta += 1
                    node.isUpdated = True
                node = node.up
                
        target_node = self.tree.search_nodes(name = str(self.target))[0]
        nodes_to_target_from_chosen = self.tree.get_distance(target_node, self.chosen_node, topology_only=True)
        return nodes_to_target_from_chosen
    
    def ChooseImages(self, setsize):
        nodes_to_chosen_image = 0
        images = []
        while len(images) < setsize:
            current_node = self.tree.get_tree_root()
            while not current_node.is_leaf():
                max_selection_criteria = 0
                for node in current_node.get_children():
                    selection_criteria = random.betavariate(node.alpha, node.beta)
                    if selection_criteria > max_selection_criteria:
                        max_selection_criteria = selection_criteria
                        max_node = node
                current_node = max_node
            #print current_node.name
            if int(current_node.name) not in images:
                if int(current_node.name) not in self.previous_images:
                    nodes_to_chosen_image += self.tree.get_distance(current_node, self.chosen_node, topology_only=True)
                    images.append(int(current_node.name))
            avnodes_to_chosen_image = nodes_to_chosen_image/setsize
        return images, avnodes_to_chosen_image
    
class OMC():
    TREE_PATH = '/home/konyushk/workspace/MCTS/Tree/'
    def __init__(self, filename, proximity_type, clustering_type, number_of_images, target):
        self.previous_images = []
        self.number_of_images = number_of_images
        self.tree = Tree(OMC.TREE_PATH + filename + str(self.number_of_images) + '_' + proximity_type + '_' + clustering_type + '.nw')
        #print len(self.tree.get_leaves())
        self.chosen_node = None
        for node in self.tree.traverse('postorder'):
            node.add_features(value = [], isUpdated = False)
            node.value.append(1)
            node.value.append(0)
        self.target = target
            
    def ModifyTree(self, images, chosen_image):
        
        for node in self.tree.traverse('postorder'):
            node.isUpdated = False
        self.previous_images = images[:]
        images.remove(chosen_image)
        not_chosen_images = images
        node = self.tree.search_nodes(name = str(chosen_image))[0]
        self.chosen_node = node
        node = node.up
        while node:
            node.value.append(1)
            node.isUpdated = True
            node = node.up
        '''for node in self.tree.traverse('postorder'):
            if node.isUpdated == False:
                    node.value.append(0)
        for node in self.tree.traverse('postorder'):
            node.isUpdated = False'''
        for image in not_chosen_images:
            node = self.tree.search_nodes(name = str(image))[0]
            node = node.up
            while node:
                if node.isUpdated == False:
                    node.value.append(0)
                    node.isUpdated = True
                node = node.up
        '''for node in self.tree.traverse('postorder'):
            if node.isUpdated == False:
                    node.value.append(1)'''
        #print self.target
        target_node = self.tree.search_nodes(name = str(self.target))[0]
        nodes_to_target_from_chosen = self.tree.get_distance(target_node, self.chosen_node, topology_only=True)
        return nodes_to_target_from_chosen
                
    def ChooseImages(self, setsize):
        nodes_to_chosen_image = 0
        images = []
        while len(images) < setsize:
            current_node = self.tree.get_tree_root()
            while not current_node.is_leaf():
                best_value = 0
                for node in current_node.get_children():
                    if numpy.mean(numpy.array(node.value)) > best_value:
                        best_value = numpy.mean(numpy.array(node.value))
                Urgency = dict()
                Fairness = dict()
                for node in current_node.get_children():
                    Urgency[node] = scipy.special.erfc((best_value-numpy.mean(numpy.array(node.value)))/(math.sqrt(2)*numpy.std(numpy.array((node.value))))) 
                for node in current_node.get_children():
                    Fairness[node] = Urgency[node]/float(sum(Urgency.values()))
                current_node = self.Sample(Fairness)
            if (int(current_node.name) not in images) and (int(current_node.name) not in self.previous_images):
                nodes_to_chosen_image += self.tree.get_distance(current_node, self.chosen_node, topology_only=True)
                images.append(int(current_node.name))
        avnodes_to_chosen_image = nodes_to_chosen_image/setsize
        return images, avnodes_to_chosen_image
    
    def Sample (self, probabilities):
                
        r = random.random()
        cumulative_sum = 0
        for node in probabilities:
            if cumulative_sum < r:
                cumulative_sum = cumulative_sum + probabilities[node]
                chosen_node = node
            else:
                break
        return chosen_node

class UCT(object):
    TREE_PATH = '/home/konyushk/workspace/MCTS/Tree/'
    C = 0.2

    def __init__(self, filename, proximity_type, clustering_type, number_of_images, target):
        self.n = 1
        self.previous_images = []
        self.number_of_images = number_of_images
        self.tree = Tree(THOMPSON.TREE_PATH + filename + str(self.number_of_images) + '_' + proximity_type + '_' + clustering_type + '.nw')
        self.chosen_node = None
        # Add alpha and beta features to the tree
        for node in self.tree.traverse('postorder'):
            node.add_features(value = [], isUpdated = False)
            node.value.append(1)
        self.target = target
        
    def ModifyTree(self, images, chosen_image):
        self.n += 1
        for node in self.tree.traverse('postorder'):
            node.isUpdated = False
        self.previous_images = images[:]
        images.remove(chosen_image)
        not_chosen_images = images
        node = self.tree.search_nodes(name = str(chosen_image))[0]
        self.chosen_node = node
        node = node.up
        while node:
            node.value.append(1)
            node.isUpdated = True
            node = node.up
        for image in not_chosen_images:
            node = self.tree.search_nodes(name = str(image))[0]
            node = node.up
            while node:
                if node.isUpdated == False:
                    node.value.append(0)
                    node.isUpdated = True
                node = node.up
        
        target_node = self.tree.search_nodes(name = str(self.target))[0]
        nodes_to_target_from_chosen = self.tree.get_distance(target_node, self.chosen_node, topology_only=True)
        return nodes_to_target_from_chosen
    
    def ChooseImages(self, setsize):
        nodes_to_chosen_image = 0
        images = []
        while len(images) < setsize:
            current_node = self.tree.get_tree_root()
            while not current_node.is_leaf():
                selection_criteria = dict()
                Fairness = dict()
                for node in current_node.get_children():
                    value = numpy.array((node.value))
                    ni = len(node.value)
                    avalue = numpy.mean(value)
                    vi = 1 / float(ni) * sum(value**2) - avalue**2 + math.sqrt(2 * math.log(self.n) / float(ni))
                    selection_criteria[node] = avalue + UCT.C * math.sqrt(math.log(self.n) / float(ni) * min(1/4,vi))
                    #print selection_criteria[node]
                for node in current_node.get_children():
                    Fairness[node] = selection_criteria[node]/float(sum(selection_criteria.values()))
                current_node = self.Sample(Fairness)
            if int(current_node.name) not in images:
                if int(current_node.name) not in self.previous_images:
                    nodes_to_chosen_image += self.tree.get_distance(current_node, self.chosen_node, topology_only=True)
                    images.append(int(current_node.name))
        avnodes_to_chosen_image = nodes_to_chosen_image/setsize
        return images, avnodes_to_chosen_image
    
    def Sample (self, probabilities):
                
        r = random.random()
        cumulative_sum = 0
        for node in probabilities:
            if cumulative_sum < r:
                cumulative_sum = cumulative_sum + probabilities[node]
                chosen_node = node
            else:
                break
        return chosen_node
    
    '''def RandomImages(self, setsize):
        images = []
        while len(images) < setsize:
            current_node = self.tree.get_tree_root()
            while not current_node.is_leaf():
                max_selection_criteria = 0
                for node in current_node.get_children():
                    selection_criteria = random.random()
                    if selection_criteria > max_selection_criteria:
                        max_selection_criteria = selection_criteria
                        max_node = node
                current_node = max_node
            if int(current_node.name) not in images:
                if int(current_node.name) not in self.previous_images:
                    images.append(int(current_node.name))
        return images'''